from typing import Any
from menu import BulletMenu
import os
from menu.helpers import Direction, clear_line, forceWrite, linebreak, move_cursor, reset_cursor, writeColor
import sys
from transformers import TrainingArguments
from enum import Enum
from myArgumentParser import myArgumentParser

def evaluation_strategy_convert(idx):
    evaluation_strategy=['no','steps','epoch']
    return evaluation_strategy[idx]

def writeColor(content, color=32, end=""):
    forceWrite(f"\u001b[{color}m{content}\u001b[0m", end)

def _ask_options(input_text, options=[], convert_value=None, default=0):
    menu = BulletMenu(input_text, options)
    result = menu.run(default_choice=default)
    return convert_value(result) if convert_value is not None else result

def _ask_field(input_text, convert_value=None, default=None, error_message=None):
    ask_again = True
    while ask_again:
        result = input(input_text)

        withdraw_length=len(list(result))
        withdraw_string=withdraw_length*'\b \b'
        sys.stdout.write('\033[F')
        print(input_text,end='')
        writeColor(f'{result}')
        print('')
        try:
            if default is not None and len(result) == 0:
                return default
            return convert_value(result) if convert_value is not None else result
        except Exception:
            if error_message is not None:
                print(error_message)

def main():
    trainingArguments=myArgumentParser((TrainingArguments))
    arg_dict={}
    cwd=os.getcwd() 
    arg_dict['output_dir'] = _ask_field(f'你想把模型放置在什么路径下？相对路径和你执行训练脚本的运行路径相关，默认是output: ',str,'output')

    arg_dict['load_best_model_at_end']=_ask_options('你想保存“最好”的checkpoint吗？',['不要','要！'],bool)
    if arg_dict['load_best_model_at_end']:
        arg_dict['metric_for_best_model']=_ask_field('⚠️ 最好的模型取决于什么指标？默认是loss。请写验证阶段返回的结果之一(最好以eval_开头，自定义的无所谓)',str,'loss')
        if arg_dict['metric_for_best_model']!='loss':
            arg_dict['greater_is_better']=_ask_options('你刚刚指定的指标大点好还是小点好？',['小点','大点'],bool)

    arg_dict['evaluation_strategy'] =_ask_options('多久验证一次？',['不验证了','每隔一定步数进行验证','每个epoch验证一次'],evaluation_strategy_convert)
    if arg_dict['evaluation_strategy']=='steps':
        arg_dict['eval_steps'] = _ask_field('每隔多少步验证一次？可以是整数，也可以是0-1区间内的浮点数（代表按总训练步数的占比进行验证） ',eval)
    arg_dict['save_strategy'] =_ask_options('多久保存一次',['不保存了','每隔一定步数保存','每个epoch保存一次'],evaluation_strategy_convert)
    if arg_dict['save_strategy']=='steps':
        if not arg_dict['load_best_model_at_end']:
            arg_dict['save_steps'] = _ask_field('每隔多少步保存？可以是整数，也可以是0-1区间内的浮点数（代表按总训练步数的占比进行验证） ',eval)
        else:
            save_times = _ask_field('每隔多少步保存一次？因为你要求保存最好的模型，所以必须是eval_step的整数倍，输入一个倍数: ',int,1)
            arg_dict['save_steps']=save_times * arg_dict['eval_steps']

    arg_dict['dataloader_drop_last'] =_ask_options('最后一个batch有可能是不完整的（剩下的样本数不够batch size），你要扔掉吗？',['扔掉','保留'],bool)
    mix_precision_dtype=_ask_options('混合精度训练？主要加速用，计算卡的半精度flops很高，大模型更适合bf16，但从A系列卡开始才支持',['不用','bf16','fp16'],int)
    if mix_precision_dtype==1:
        arg_dict['bf16']=True
    elif mix_precision_dtype==2:
        arg_dict['fp16']=True

    arg_dict['gradient_accumulation_steps'] =_ask_field('梯度累计多少步？ 🤗训练时显示的总step约等于总样本数/batch_size/梯度累计步数: ',eval)
    arg_dict['auto_find_batch_size']=_ask_options('🤗可以帮你自动决定batch_size，这个特性需要pip install accelerate',['我不','好哇！'],bool)
    if not arg_dict['auto_find_batch_size']:
        arg_dict['per_device_train_batch_size'] =_ask_field('不要自动决定？那每张卡上batch_size多少？',eval)
        
    arg_dict['gradient_checkpointing']=_ask_options('检查点激活，牺牲训练速度换来显存减少',['不要','要！'],bool)
    arg_dict['remove_unused_columns']  =_ask_options('⚠️ 帮你把模型forward不需要的键值对从collate返回值里删了，你必须很清楚这是在干什么',['保留','扔了'],bool,0)

    arg_dict['lr_scheduler_type ']  =_ask_options('⚠️ 学习率调节器的类型',['linear','cosine','cosine_with_restarts','polynomial','constant'],bool,0)
    

    args = trainingArguments.parse_dict(arg_dict)

if __name__=='__main__':
    main()